import torch
import torch.nn as nn
from torch.nn import Transformer
import jieba
from collections import Counter
from torch.utils.data import Dataset, DataLoader
import re
import math

input_sentence = ["你好啊", "你多大了", "我想去吃午饭", '你给我讲个故事']
output_sentence = ["我很好", "我今年18岁了", "正好，我们一起去", '从前有座山，山上有座庙，庙里有两个和尚']

n_step1 = max([len(i) for i in input_sentence])
n_step2 = max([len(i) for i in output_sentence])
n_step = max(n_step1, n_step2) + 2


# 建立字典
def buildDic(texts):
    resultAll = []
    for text in texts:
        if isinstance(text, list):
            for t in text:
                t = re.sub(r'[^\u4e00-\u9fa5a-zA-Z0-9，。！？]', '', t)
                result = jieba.cut(t)
                resultAll.extend(result)
    count = Counter(resultAll)
    word_vacab = count.most_common(9999)
    word_vacab = [w for w, i in word_vacab]
    word_vacab.insert(0, "<PAD>")
    word_vacab.insert(1, "<UNK>")
    word_vacab.insert(2, "<SOS>")
    word_vacab.insert(3, "<EOS>")

    word2idx = {word: idx for idx, word in enumerate(word_vacab)}
    idx2word = {idx: word for idx, word in enumerate(word_vacab)}

    return word_vacab, word2idx, idx2word


# 字典建立
word_vacab, word2idx, idx2word = buildDic([input_sentence, output_sentence])
print(word2idx)


def preprocess_sentence(sentence, word2idx, n_step):
    # 字符过滤
    sentence = re.sub(r'[^\u4e00-\u9fa5a-zA-Z0-9，。！？]', '', sentence)
    # 分词
    words = jieba.cut(sentence)
    # 索引映射
    array = [2] + [word2idx.get(w, word2idx["<UNK>"]) for w in words] + [3]
    # 序列填充
    if len(array) < n_step:
        array += [0] * (n_step - len(array))
    return array


def make_data(input_sentence, output_sentence):
    input_array, output_array = [], []
    for input_sent, output_sent in zip(input_sentence, output_sentence):
        input_array.append(preprocess_sentence(input_sent, word2idx, n_step))
        output_array.append(preprocess_sentence(output_sent, word2idx, n_step))

    return torch.Tensor(input_array).long(), torch.Tensor(output_array).long()


input_tensor, output_tensor = make_data(input_sentence, output_sentence)


class MyDataSet(Dataset):
    def __init__(self, input_tensor, output_tensor):
        self.input_tensor = input_tensor
        self.output_tensor = output_tensor

    def __len__(self):
        return len(self.input_tensor)

    def __getitem__(self, idx):
        return self.input_tensor[idx], self.output_tensor[idx]


# 超参数设置
VOCAB_SIZE = len(word_vacab)  # 词的最大维度 也就是词典一共可以识别多少组词
BATCH_SIZE = 1  # 数据集 每批次的样本数量
EMBED_SIZE = 128  # 词嵌入的维度
N_HEAD = 4  # 多头 头数
NUM_ENCODER_LAYERS = 3  # Encoder层数
NUM_DECODER_LAYERS = 3  # Decoder层数
FFN_HID_DIM = 256  # 前馈隐藏层维度

loader = DataLoader(MyDataSet(input_tensor, output_tensor), BATCH_SIZE, shuffle=True)


class PositionalEncoding(nn.Module):
    def __init__(self, d_model, dropout=0.1, max_len=5000):
        super().__init__()
        self.dropout = nn.Dropout(p=dropout)

        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0)
        self.register_buffer('pe', pe)

    def forward(self, x):
        # (1, max_len, d_model)
        x = x + self.pe[:, :x.size(1), :]
        return self.dropout(x)


class MyTransformer(nn.Module):
    def __init__(self):
        super(MyTransformer, self).__init__()
        self.embedding = nn.Embedding(VOCAB_SIZE, EMBED_SIZE)
        self.positional_encoding = PositionalEncoding(EMBED_SIZE)
        self.transformer = Transformer(
            d_model=EMBED_SIZE,
            nhead=N_HEAD,
            num_decoder_layers=NUM_DECODER_LAYERS,
            num_encoder_layers=NUM_ENCODER_LAYERS,
            dim_feedforward=FFN_HID_DIM,
            dropout=0.1,
            batch_first=True
        )
        self.fc = nn.Linear(EMBED_SIZE, VOCAB_SIZE)

    def forward(self, src, tgt):
        src = self.positional_encoding(self.embedding(src))
        tgt = self.positional_encoding(self.embedding(tgt))

        # 生成mask
        src_mask = self.transformer.generate_square_subsequent_mask(src.size(1)).to(src.device)
        tgt_mask = self.transformer.generate_square_subsequent_mask(tgt.size(1)).to(tgt.device)

        # Transformer前向传播
        # output (batch_size, tgt_seq_len, d_model)
        output = self.transformer(
            src,
            tgt,
            src_mask=src_mask,
            tgt_mask=tgt_mask,
            memory_mask=None
        )
        return self.fc(output)


device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = MyTransformer().to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
criterion = nn.CrossEntropyLoss(ignore_index=0)  # 假设0是<pad>标记


def train(model, src, tgt):
    optimizer.zero_grad()

    # 将输入和目标移动到设备
    src = src.to(device)
    tgt = tgt.to(device)

    # 前向传播
    output = model(src, tgt[:, :-1])  # 解码器输入使用左移的target

    # 计算损失
    loss = criterion(output.view(-1, VOCAB_SIZE), tgt[:, 1:].reshape(-1))
    loss.backward()
    optimizer.step()
    return loss.item()


# 训练示例（实际需多轮迭代）
for epoch in range(100):
    model.train()
    total_loss = 0

    for batch_x, batch_y in loader:
        total_loss += train(model, batch_x, batch_y)

    if (epoch + 1) % 10 == 0:
        print(f'Epoch {epoch + 1}, Loss: {total_loss:.4f}')


def predictSay(text):
    model.eval()
    myinput, _ = make_data([text], [""])

    with torch.no_grad():
        tgt = torch.tensor([[word2idx["<SOS>"]]]).to(device)
        decoded = []

        # 逐步生成输出序列
        for _ in range(n_step):
            output = model(myinput.to(device), tgt)  # 前向传播
            next_word_idx = output[:, -1, :].argmax(dim=-1).item()  # 获取下一个单词的索引
            decoded.append(idx2word[next_word_idx])  # 将索引转换为单词并添加到结果中
            # 如果生成了 <EOS>，停止生成
            if next_word_idx == word2idx["<EOS>"]:
                break

            # 更新目标序列
            tgt = torch.cat([tgt, torch.tensor([[next_word_idx]], dtype=torch.long).to(device)], dim=1)

    return ''.join(decoded).replace("<SOS>", "").replace("<PAD>", "").replace("<EOS>", "")


print(predictSay("你多大了"))

# torch.save(model, "pytorch-transformer.pth")  # 保存模型参数
# 转换为 TorchScript 并保存
scripted_model = torch.jit.script(model)
scripted_model.save("pytorch-transformer.pt")